"""
requirements:
kaggle          从kaggle网站上下载数据
datasets        Transformer输入需要的数据结构
transformers    分词模型加载，分词器加载
pytorch         模型加载
tensorflow      模型加载
chardet         The Universal Character Encoding Detector
sentencepiece   Python wrapper for SentencePiece. This API will offer the encoding, decoding and training of Sentencepiece.

"""

import os
from pathlib import Path

# 设置是否是 kaggle 环境
is_kaggle = os.environ.get('KAGGLE_KERNEL_RUN_TYPE', '')

# kaggle 账号
credit = '{"username": "owlion", "key": "e0f53f676521d65a4af2867168829452"}'

# 创建kaggle配置文件
credit_path = Path("~/.kaggle/kaggle.json").expanduser()
if not credit_path.exists():
    credit_path.parent.mkdir(exist_ok=True)
    credit_path.write_text(credit)
    credit_path.chmod(0o600)

"""
1. 数据准备阶段
"""
# 下载测试数据
path = Path('../data/us-patent-phrase-to-phrase-matching')
if not is_kaggle and not path.exists():
    import zipfile, kaggle

    kaggle.api.competition_download_cli(str(path))
    zipfile.ZipFile(f'{path}.zip').extractall(path)

"""
2. 数据整理阶段
"""
import pandas as pd

# -- 读取数据
df = pd.read_csv('../data/us-patent-phrase-to-phrase-matching/train.csv')
print(df.describe(include='object'))
print(df)

# -- 创建输入数据格式
df['input'] = 'TEXT1: ' + df.context + '; TEXT2: ' + df.target + '; ANC1: ' + df.anchor
print(df.input.head())

"""
3. 数据分词阶段
"""
# -- 数据分词
from datasets import Dataset, DatasetDict

# Transformers 使用 Dataset 对象存储数据，将 DataFrame 转成 Dataset
dataset = Dataset.from_pandas(df)
print(dataset)

'''
But we can't pass the texts directly into a model. 
A deep learning model expects numbers as inputs, 
not English sentences! So we need to do two things:

Tokenization: Split each text up into words (or actually, as we'll see, into tokens)
Numericalization: Convert each word (or token) into a number.
'''
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# 使用模型名字在程序第一次运行时下载模型到本地
model_name = 'microsoft/deberta-v3-small'
# 提前下载好模型，模型加载时提供模型的路径
model_path = "/Users/ws/Workspace/llm/models/deberta-v3-small"
# 通过路径加载模型, 使用use_fast会报警告
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
example = "G'day folks, I'm Jeremy from fast.ai!"
print(f"测试分词数据：{example}")
print(f"测试分词结果：{tokenizer.tokenize(example)}")
print(f"测试分词对象：{tokenizer(example)}\n")


# 定义一个函数，对传入对象中的 input 字段进行分词
def tokenize_input(x):
    return tokenizer(x["input"])


# 通过并行方式批量对数据集中的 input 字段进行分词
# 分词后，会在原有数据集中增加一个 input_ids 字段
# What are those IDs and where do they come from?
# The secret is that there's a list called vocab in the tokenizer which contains a unique integer
# for every possible token string. We can look them up like this, for instance to find the token for the word "of":
# 注意 ▁of 前缀不是下划线，而是 space；而且 ▁of 和 of 在词库中的 id 是不一样的
token_dataset = dataset.map(tokenize_input, batched=True)
print(f"token_dataset type is {type(token_dataset)}\n")
print(f"token_dataset[0]['input'] : {token_dataset[0]["input"]}\n")
print(f"token_dataset[0]['input_ids'] : {token_dataset[0]["input_ids"]}\n")
print(f"_of 在词库中的id: {tokenizer.vocab['▁of']}")
print(f"of 在词库中的id: {tokenizer.vocab['of']}")

# Transformers always assumes that your labels has the column name labels, but in our dataset it's currently score.
# Therefore, we need to rename it
# 后续采用的是 AutoModelForSequenceClassification 模型，用于进行序列分类任务。
# 它的作用是根据提供的预训练模型名称（如BERT、RoBERTa等）自动加载对应的模型，并添加一个额外的分类头部，以适应特定的序列分类任务。
# 分类一般应该是输出字符型的类别名字，但这里是用的数值类型的 score 作为类别
token_dataset = token_dataset.rename_columns({"score": "labels"})

"""
4. 创建测试集和验证集数据
"""

# Transformers uses a DatasetDict for holding your training and validation sets.
# To create one that contains 25% of our data for the validation set, and 75% for the training set

# dataset_dict 中包含 train 训练集，和 test 验证集（其实用validate命名更合适）
dataset_dict = token_dataset.train_test_split(0.25, seed=42)
print(dataset_dict)

# 准备好最后的测试集数据
df_test = pd.read_csv("../data/us-patent-phrase-to-phrase-matching/test.csv")
df_test.describe()
df_test['input'] = 'TEXT1: ' + df_test.context + '; TEXT2: ' + df_test.target + '; ANC1: ' + df_test.anchor
dataset_test = Dataset.from_pandas(df_test).map(tokenize_input, batched=True)

"""
4. 创建测试集和验证集数据
"""

import numpy as np
from sklearn.metrics import accuracy_score, f1_score

def corr(x, y):
    # 用于计算给定数组的皮尔逊相关系数矩阵。它的作用是计算输入数组中各列之间的皮尔逊相关系数，并返回一个相关系数矩阵，
    print(f"corr({x}, {y}) = {np.corrcoef(x, y)[0][1]}")
    return np.corrcoef(x, y)[0][1]


def corr_d(pred):
    """
    接收一个 transformers.trainer_utils.EvalPrediction 对象，包含两个属性
        predictions (:obj:`np.ndarray`): Predictions of the model. 本轮结果的预测值
        label_ids (:obj:`np.ndarray`): Targets to be matched. 训练集中的真实值
    这里通过corr函数计算两组值之间的皮尔逊相关系数，如果相关系数变大，说明模型训练在向好的方向发展
    返回的内容将放入到输出信息中，构成训练过程状态参数 matrix 中的一项，例如：
        {'eval_loss': 0.023677006363868713, 'eval_皮尔逊相关系数': 0.819147976220058, 'eval_runtime': 8.7089,
        'eval_samples_per_second': 1047.089, 'eval_steps_per_second': 4.134, 'epoch': 4.0}
    """
    return {'皮尔逊相关系数': corr(*pred)}


"""
5. 训练模型
"""
from transformers import TrainingArguments, Trainer

# 选择批处理数量，以便适合GPU数量
batche_size = 256
# 迭代次数
epochs = 3
# 最重要的超参数是学习率，只能使用试错法。找到不会导致训练失败的最大值。
learning_rate = 8e-5

"""
参数说明：
    'outputs': 指定训练过程中输出文件的目录。
    learning_rate=8e-5: 学习率，控制模型在每次更新时参数的调整大小。较小的学习率通常使得模型学习更加稳定，但可能需要更多的训练时间。
    warmup_ratio=0.1: 学习率预热比例。预热是指在训练开始时逐渐增加学习率的过程，使得模型更容易收敛到合适的参数。这里设置为总训练步数的10%。
    lr_scheduler_type='cosine': 学习率调度器的类型。这里使用余弦退火调度器，它会在训练过程中逐渐降低学习率，有助于提高模型性能。
    evaluation_strategy='epoch': 评估策略。指定在何时评估模型性能，这里设置为每个 epoch 结束时评估一次。
    per_device_train_batch_size=128: 每个设备上的训练批量大小。控制每次模型更新时输入的样本数量，影响训练速度和内存需求。
    per_device_eval_batch_size=128 * 2: 每个设备上的评估批量大小。控制每次评估时输入的样本数量，同样影响速度和内存需求。
    num_train_epochs=4: 训练的总 epoch 数量。指定模型在整个训练数据集上迭代的次数。
    weight_decay=0.01: 权重衰减（L2正则化）系数。用于控制模型参数的大小，防止过拟合。
    report_to=None: 报告日志的目标。这里设置为 None，表示不生成任何训练日志。
"""
args = TrainingArguments(
    'outputs',
    learning_rate=learning_rate,
    warmup_ratio=0.1,
    lr_scheduler_type='cosine',
    # fp16=True,
    evaluation_strategy='epoch',
    per_device_train_batch_size=batche_size,
    per_device_eval_batch_size=batche_size * 2,
    num_train_epochs=4,
    weight_decay=0.01,
    report_to=None
)

"""
参数说明：
    model_name: 指定要使用的预训练模型的名称或路径。
    num_labels=1: 序列分类任务的标签数量。在这种情况下，该模型是一个二分类模型，因此将标签数量设置为1。
                  实际情况是输出是（0，1）之间的一个数值，表示专利数据集中 TEXT2 和 ANC1 字段在 TEXT1 上下文环境下的相似度
"""
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    num_labels=1
)

"""
参数说明：
    compute_metrics：用于指定一个函数，用于计算训练或评估过程中的指标（metrics）。
    具体来说，这个参数允许你自定义一个函数，该函数将接收模型预测的输出和真实标签作为输入，并返回一个字典，包含你感兴趣的任何指标。
"""
trainer = Trainer(
    model, args,
    train_dataset=dataset_dict['train'],
    eval_dataset=dataset_dict['test'],
    tokenizer=tokenizer,
    compute_metrics=corr_d
)

trainer.train()

predicts = trainer.predict(dataset_test).predictions.astype(float)
predicts = np.clip(predicts, 0, 1)
print(predicts)